const std = @import("std");
const json = @import("std").json;

const Allocator = std.mem.Allocator;

fn cleanNotebook(allocator: Allocator, input_path: []const u8, output_path: []const u8) !void {
    const file = try std.fs.cwd().openFile(input_path, .{});
    defer file.close();

    const stat = try file.stat();
    const content = try file.readToEndAlloc(allocator, stat.size);
    defer allocator.free(content);

    var parsed = try json.parseFromSlice(std.json.Value, allocator, content, .{});
    defer parsed.deinit();

    const root = parsed.value;

    if (root != .Object) return error.InvalidNotebook;

    const cells = root.Object.get("cells") orelse return error.NoCells;
    if (cells != .Array) return error.InvalidCells;

    for (cells.Array.items) |*cell| {
        if (cell.* != .Object) continue;
        const cell_obj = &cell.*.Object;

        // 清理 outputs
        if (cell_obj.get("outputs")) |outputs| {
            if (outputs == .Array) {
                outputs.Array.items.len = 0;
            }
        }

        // 清理 execution_count
        if (cell_obj.get("execution_count")) |exec_count| {
            _ = exec_count;
            cell_obj.put("execution_count", std.json.Value{ .Null = {} }) catch continue;
        }
    }

    const output_file = try std.fs.cwd().createFile(output_path, .{});
    defer output_file.close();

    const out_stream = output_file.writer();
    try std.json.stringify(root, .{ .indent = .{ .Space = 2 } }, out_stream);
}

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    const allocator = gpa.allocator();

    const args = try std.process.argsAlloc(allocator);
    defer std.process.argsFree(allocator, args);

    if (args.len < 2) {
        std.debug.print("Usage: {s} <notebook.ipynb> [output.ipynb]\n", .{args[0]});
        return;
    }

    const input = args[1];
    const output = if (args.len >= 3) args[2] else input;

    cleanNotebook(allocator, input, output) catch |err| {
        std.debug.print("Error: {}\n", .{err});
        std.process.exit(1);
    };

    std.debug.print("✅ Cleaned notebook saved to: {s}\n", .{output});
}
